import numpy as np
from random import uniform


class record(object):
	def __init__(self):
		self.summ = 0.		# summation f
		self.sumSquare = 0.	# summation f^2
		self.count = 0

	def Clear(self):
		self.summ = 0.
		self.sumSquare = 0.
		self.count = 0

	def Add(self, value):
		self.count += 1
		self.summ += value
		self.sumSquare += value*value

	def GetAvg(self):
		return self.summ/float(self.count)	# return <f>

	def GetAvgSquare(self):
		return self.sumSquare/float(self.count)	# return <f^2>
	# end of record


SPINUP = 1.	# define Spin up = 1.
SPINDOWN = -1.	# define Spin down = -1.

class twoD_IsingModel(object):
	def __init__(self, dim = (1, 1), k = 1., J = 1.):
		self.energy = 0.	# total energy of lattice
		self.sumSpins = 0.	# total magnetic field of lattice
		self.k = k		# Bolzmann constant
		self.J = J		# spin energy
		self.nx = dim[0]
		self.ny = dim[1]
		self.lattice = np.zeros((self.nx, self.ny))	# create lattice
		self.area = float(self.nx * self.ny)
		self.E_Record = record()	# record energy of lattice
		self.sumSpins_Record = record()	# record magnetism of lattice

	def SetTemperatureField(self, J = 1., T = 10., H = 0.):
		self.H = H	# external magnetic field
		self.T = T	# temperature
		self.beta = 1./self.T/self.k

	def Thermalize(self):
		for i in xrange(self.nx):
			for j in xrange(self.ny):
				self.lattice[i,j] = SPINUP if int(uniform(0, 2))==1 else SPINDOWN


	def CalcEnergy(self):
		tmpE = 0.
		for i in xrange(self.nx):
			for j in xrange(self.ny):
				tmpE += -self.J \
				* ( self.lattice[i,j]*self.lattice[i,j-1] \
				+ self.lattice[i,j]*self.lattice[i-1,j] ) \
				- self.H*self.lattice[i,j]
		self.energy = tmpE

	def CalcSumSpins(self):
		self.sumSpins = self.lattice.sum()

	def FlipSpin(self, ix, iy):
		self.lattice[ix,iy] = SPINDOWN if self.lattice[ix,iy] == SPINUP else SPINUP

	def GetNeighborEnergy(self, ix, iy):
		f_ix = (ix+1) % self.nx
		b_ix = ix-1
		f_iy = (iy+1) % self.ny
		b_iy = iy-1
		E = -( self.lattice[f_ix,iy] + self.lattice[b_ix,iy] \
		+ self.lattice[ix,f_iy] + self.lattice[ix,b_iy] ) \
		* self.lattice[ix,iy] * self.J \
		- self.H * self.lattice[ix,iy]
		return E

	def Prob(self, deltaE):
		if np.e**(-deltaE*self.beta) > uniform(0, 1):
			return True
		else:
			return False

	def LocalUpdate(self, x, y):

		oldE = self.GetNeighborEnergy(x, y)
		oldSpin = self.lattice[x,y]

		self.FlipSpin(x, y)

		newE = self.GetNeighborEnergy(x, y)
		newSpin = self.lattice[x,y]
		deltaE = newE - oldE
		deltaSpin = newSpin - oldSpin

		if (deltaE <= 0) or self.Prob(deltaE):
			self.energy += deltaE
			self.sumSpins += deltaSpin
		else:
			self.FlipSpin(x, y)

	def Start(self, N, thermalStep):
		self.E_Record.Clear()
		self.sumSpins_Record.Clear()

		self.CalcEnergy()
		self.CalcSumSpins()

		for i in xrange(N):
			x = int(uniform(0, self.nx))
			y = int(uniform(0, self.ny))
			self.LocalUpdate(x, y)
			if i > thermalStep:
				self.sumSpins_Record.Add(self.GetSumSpins())
				self.E_Record.Add(self.GetEnergy())
	
	def GetEnergy(self):
		return self.energy

	def GetSumSpins(self):
		return self.sumSpins

	def GetAvgEnergy(self):
		return self.E_Record.GetAvg()/self.area

	def GetAvgMagnetism(self):
		return self.sumSpins_Record.GetAvg()/self.area

	def GetAvgSquarMagnetism(self):
		return self.sumSpins_Record.GetAvgSquare()/self.area**2.

	def GetSpecificHeat(self):
		return ( self.E_Record.GetAvgSquare()/self.area**2. \
		- (self.E_Record.GetAvg()/self.area)**2. ) \
		* self.beta/self.T

	def GetMagneticSusceptibility(self):
		return ( self.sumSpins_Record.GetAvgSquare()/self.area**2. \
		- (self.sumSpins_Record.GetAvg()/self.area)**2. ) \
		* self.beta
	# end of twoD_IsingModel


def SaveData(savePath="", data=[], labels=""):
	fp = open(savePath, 'w')
	fp.write(labels + '\n')
	for column in data:
		for rowItem in column:
			fp.write(str(rowItem)+'\t')
		fp.write('\n')
	fp.close()
	# end of SaveData


def LoopT_IsingModel( model, TList ):
	avgE_List = []		#<E>
	avgM_List = []		#<M>
	avgSquarM_List = []	# <M^2>
	MS_List = []	# Magnetic suscep
	SH_List = []	# Specific heat

	model.Thermalize()
	for i in TList:
		print "Temperature", i
		model.SetTemperatureField( T = i, H = 0. )
		model.Start( iterN, thermalStep )

		avgE_List.append( model.GetAvgEnergy() )
		avgM_List.append( model.GetAvgMagnetism() )
		avgSquarM_List.append( model.GetAvgSquarMagnetism() )
		MS_List.append( model.GetMagneticSusceptibility() )
		SH_List.append( model.GetSpecificHeat() )

	dataLabel = 'temperature\tM_Square\tSuscep\tSpecificHeat\t'

	return dataLabel, \
		[ temperatureList, \
		avgSquarM_List, \
		MS_List, \
		SH_List \
		]
	
if __name__ == '__main__':
	import time
	L = 20		# lattice size
	dim = ( L, L )	# dim of lattice is L*L

	iterNumPerSite = 10
	iterN = L**2 * iterNumPerSite	# iter times
	thermalStep = int( iterN * 0.3 )	# reach thermalize after thermalStep
	temperatureList = np.linspace( 5., 1., 41 )	# cooling down

	tic = time.time()

	model = twoD_IsingModel( dim )
	dataLabel, tmpData = LoopT_IsingModel( model, temperatureList )

	toc = time.time()
	print "Time ", toc-tic

	allData = np.array( tmpData ).transpose()
	savePath = "2D_IsingModelResult.txt"
	SaveData( savePath = savePath, data = allData, labels = dataLabel )

	## reshape data for gnuplot
	#avgEList = []
	#avgMList = []
	#avgSquarMList = []
	#SHList = []
	#MSList = []
	#for i in range(len(temperatureList)):
	#	avgEList.append([temperatureList[i], avgE_List[i]])
	#	avgMList.append([temperatureList[i], avgM_List[i]])
	#	avgSquarMList.append([temperatureList[i], avgSquarM_List[i]])
	#	SHList.append([temperatureList[i], SH_List[i]])
	#	MSList.append([temperatureList[i], MS_List[i]])


	#import Gnuplot
	#g = Gnuplot.Gnuplot()
	#g.xlabel("Temperature (K)")
	#g.ylabel("<E>")
	#g.plot(avgEList)

	#g2 = Gnuplot.Gnuplot()
	#g2.xlabel("Temperature (K)")
	#g2.ylabel("<M>")
	#g2.plot(avgMList)

	#g3 = Gnuplot.Gnuplot()
	#g3.xlabel("Temperature (K)")
	#g3.ylabel("Cv")
	#g3.plot(SHList)

	#g4 = Gnuplot.Gnuplot()
	#g4.xlabel("Temperature (K)")
	#g4.ylabel("X")
	#g4.plot(MSList)
	#raw_input()
